import numpy as np

# 定义游戏状态类
class GameState:
    def __init__(self, sun_count, zombies, plants, game_status, wave, lawnmowers, time_remaining):
        self.sun_count = sun_count
        self.zombies = zombies
        self.plants = plants
        self.game_status = game_status
        self.wave = wave
        self.lawnmowers = lawnmowers
        self.time_remaining = time_remaining

    def get_state_vector(self):
        state = [self.sun_count, self.wave, self.time_remaining]
        for zombie in self.zombies:
            state.extend([zombie['row'], zombie['position'], zombie['health']])
        for plant in self.plants:
            state.extend([plant['row'], plant['position'], plant['cooldown']])
        for lawnmower in self.lawnmowers:
            state.append(int(lawnmower['available']))
        return np.array(state)

# 定义 Q - learning 智能体类
class QLearningAgent:
    def __init__(self, state_size, action_size, learning_rate=0.1, discount_factor=0.9, exploration_rate=0.1):
        self.state_size = state_size
        self.action_size = action_size
        self.learning_rate = learning_rate
        self.discount_factor = discount_factor
        self.exploration_rate = exploration_rate
        self.q_table = np.zeros((state_size, action_size))

    def choose_action(self, state):
        if np.random.uniform(0, 1) < self.exploration_rate:
            action = np.random.choice(self.action_size)
        else:
            state_index = hash(tuple(state)) % self.state_size
            action = np.argmax(self.q_table[state_index])
        return action

    def load_model(self, filename):
        self.q_table = np.load(filename)

# 定义动作空间和植物种类
ROWS = 5
PLANT_TYPES = ["Sunflower", "Peashooter"]
ACTION_SIZE = ROWS * len(PLANT_TYPES)

state_size = 1000

# 创建 Q - learning 智能体
agent = QLearningAgent(state_size, ACTION_SIZE)

# 加载保存的模型
model_filename = "D:\Project\Pycharm Project\pythonProject4\models\direction_model.npy"
agent.load_model(model_filename)

# 模拟一个新的游戏状态
new_game_data = {
    "sun_count": 200,
    "zombies": [
        {"type": "normal", "row": 2, "position": 350, "health": 120}
    ],
    "plants": [
        {"type": "Peashooter", "row": 2, "position": 150, "cooldown": 3}
    ],
    "game_status": "playing",
    "wave": 4,
    "lawnmowers": [
        {"row": 1, "available": True},
        {"row": 2, "available": True},
        {"row": 3, "available": False},
        {"row": 4, "available": True},
        {"row": 5, "available": True}
    ],
    "time_remaining": 25
}
new_game_state = GameState(new_game_data["sun_count"], new_game_data["zombies"], new_game_data["plants"],
                           new_game_data["game_status"], new_game_data["wave"], new_game_data["lawnmowers"],
                           new_game_data["time_remaining"])
new_state_vector = new_game_state.get_state_vector()

# 解析动作到种植操作
def action_to_planting(action):
    row = action % ROWS + 1
    plant_type_index = action // ROWS
    plant_type = PLANT_TYPES[plant_type_index]
    return row, plant_type

# 选择动作
action = agent.choose_action(new_state_vector)

# 解析动作到种植操作
row, plant_type = action_to_planting(action)

print(f"建议在第 {row} 行种植 {plant_type}")